package com.flink.hbase;

import org.apache.flink.api.common.state.BroadcastState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.state.ReadOnlyBroadcastState;
import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.streaming.api.functions.co.BroadcastProcessFunction;
import org.apache.flink.util.Collector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;
import java.util.Map;

/**
 * 用户数据丰富化处理函数
 * 使用广播流中的维度数据来丰富主流中的用户数据
 */
public class UserEnrichmentFunction extends BroadcastProcessFunction<User, DimensionData, EnrichedUser> {

    private static final Logger LOG = LoggerFactory.getLogger(UserEnrichmentFunction.class);

    // 广播状态描述符 - 存储维度数据
    public static final MapStateDescriptor<String, DimensionData> DIMENSION_STATE_DESCRIPTOR =
            new MapStateDescriptor<>(
                    "dimension-state",
                    BasicTypeInfo.STRING_TYPE_INFO,
                    TypeInformation.of(DimensionData.class)
            );

    @Override
    public void processElement(User user, 
                             ReadOnlyContext readOnlyContext, 
                             Collector<EnrichedUser> collector) throws Exception {
        
        if (user == null) {
            LOG.warn("接收到空的用户数据");
            return;
        }

        LOG.debug("处理用户数据: {}", user.getUserId());

        // 获取广播状态
        ReadOnlyBroadcastState<String, DimensionData> broadcastState = 
                readOnlyContext.getBroadcastState(DIMENSION_STATE_DESCRIPTOR);

        // 创建丰富化的用户对象
        EnrichedUser enrichedUser = new EnrichedUser(user);

        // 尝试从广播状态中获取相关的维度数据
        try {
            // 构建维度查找键
            Map<String, String> lookupKeys = buildLookupKeys(user);
            
            // 应用维度数据进行丰富化
            for (Map.Entry<String, String> entry : lookupKeys.entrySet()) {
                String category = entry.getKey();
                String lookupKey = entry.getValue();
                
                DimensionData dimensionData = broadcastState.get(lookupKey);
                if (dimensionData != null && dimensionData.getIsActive()) {
                    enrichedUser.enrichWithDimension(dimensionData);
                    LOG.debug("用户 {} 使用维度 {} 进行丰富化: {}", 
                             user.getUserId(), category, dimensionData.getDimensionValue());
                }
            }

            // 设置丰富化状态标志
            if (!enrichedUser.isFullyEnriched()) {
                LOG.warn("用户 {} 未完全丰富化，缺少部分维度信息", user.getUserId());
            }

            // 输出丰富化后的用户数据
            collector.collect(enrichedUser);
            
        } catch (Exception e) {
            LOG.error("处理用户数据丰富化时发生异常: userId={}", user.getUserId(), e);
            
            // 异常情况下仍然输出基础用户数据
            collector.collect(enrichedUser);
        }
    }

    @Override
    public void processBroadcastElement(DimensionData dimensionData, 
                                      Context context, 
                                      Collector<EnrichedUser> collector) throws Exception {
        
        if (dimensionData == null) {
            LOG.warn("接收到空的维度数据");
            return;
        }

        LOG.info("处理广播维度数据: category={}, key={}, value={}", 
                 dimensionData.getCategory(), 
                 dimensionData.getDimensionKey(), 
                 dimensionData.getDimensionValue());

        // 获取广播状态
        BroadcastState<String, DimensionData> broadcastState = 
                context.getBroadcastState(DIMENSION_STATE_DESCRIPTOR);

        try {
            // 构建状态存储键
            String stateKey = buildStateKey(dimensionData);
            
            if (dimensionData.getIsActive()) {
                // 更新或添加维度数据
                broadcastState.put(stateKey, dimensionData);
                LOG.debug("更新广播状态: key={}, data={}", stateKey, dimensionData);
            } else {
                // 删除非活跃的维度数据
                broadcastState.remove(stateKey);
                LOG.debug("删除非活跃维度数据: key={}", stateKey);
            }

            // 定期清理过期数据
            cleanupExpiredData(broadcastState, dimensionData.getUpdateTime());
            
        } catch (Exception e) {
            LOG.error("处理广播维度数据时发生异常: {}", dimensionData, e);
        }
    }

    /**
     * 构建维度数据的状态存储键
     */
    private String buildStateKey(DimensionData dimensionData) {
        return dimensionData.getCategory() + ":" + dimensionData.getDimensionKey();
    }

    /**
     * 构建用户数据的维度查找键
     */
    private Map<String, String> buildLookupKeys(User user) {
        Map<String, String> lookupKeys = new HashMap<>();
        
        // 根据用户 ID 查找用户分类
        lookupKeys.put("user_category", "user_category:" + user.getUserId());
        
        // 根据用户 ID 查找用户等级
        lookupKeys.put("user_level", "user_level:" + user.getUserId());
        
        // 根据年龄段查找地区信息（示例逻辑）
        String ageGroup = getAgeGroup(user.getAge());
        lookupKeys.put("region", "region:" + ageGroup);
        
        // 根据性别查找 VIP 状态（示例逻辑）
        lookupKeys.put("vip_status", "vip_status:" + user.getGender());
        
        return lookupKeys;
    }

    /**
     * 根据年龄获取年龄组
     */
    private String getAgeGroup(Integer age) {
        if (age == null) {
            return "unknown";
        }
        
        if (age < 18) {
            return "young";
        } else if (age < 35) {
            return "adult";
        } else if (age < 50) {
            return "middle";
        } else {
            return "senior";
        }
    }

    /**
     * 清理过期的维度数据
     */
    private void cleanupExpiredData(BroadcastState<String, DimensionData> broadcastState, 
                                   Long currentTime) throws Exception {
        
        // 设置过期时间为 24 小时
        long expirationThreshold = currentTime - (24 * 60 * 60 * 1000);
        
        // 遍历状态中的所有数据，移除过期数据
        // 注意：这个操作可能比较昂贵，实际生产中可以考虑更高效的清理策略
        Iterable<Map.Entry<String, DimensionData>> entries = broadcastState.entries();
        for (Map.Entry<String, DimensionData> entry : entries) {
            DimensionData data = entry.getValue();
            if (data.getUpdateTime() < expirationThreshold) {
                broadcastState.remove(entry.getKey());
                LOG.debug("清理过期维度数据: key={}", entry.getKey());
            }
        }
    }

    /**
     * 获取广播状态中的统计信息（用于监控）
     */
    public static class StateStatistics {
        private final int totalDimensions;
        private final Map<String, Integer> categoryCount;
        private final long lastUpdateTime;

        public StateStatistics(int totalDimensions, Map<String, Integer> categoryCount, long lastUpdateTime) {
            this.totalDimensions = totalDimensions;
            this.categoryCount = categoryCount;
            this.lastUpdateTime = lastUpdateTime;
        }

        // Getters
        public int getTotalDimensions() { return totalDimensions; }
        public Map<String, Integer> getCategoryCount() { return categoryCount; }
        public long getLastUpdateTime() { return lastUpdateTime; }

        @Override
        public String toString() {
            return "StateStatistics{" +
                   "totalDimensions=" + totalDimensions +
                   ", categoryCount=" + categoryCount +
                   ", lastUpdateTime=" + lastUpdateTime +
                   '}';
        }
    }

    /**
     * 计算广播状态统计信息
     */
    private StateStatistics calculateStateStatistics(BroadcastState<String, DimensionData> broadcastState) 
            throws Exception {
        
        Map<String, Integer> categoryCount = new HashMap<>();
        int totalCount = 0;
        long maxUpdateTime = 0;

        Iterable<Map.Entry<String, DimensionData>> entries = broadcastState.entries();
        for (Map.Entry<String, DimensionData> entry : entries) {
            DimensionData data = entry.getValue();
            
            // 统计分类数量
            categoryCount.merge(data.getCategory(), 1, Integer::sum);
            totalCount++;
            
            // 记录最新更新时间
            maxUpdateTime = Math.max(maxUpdateTime, data.getUpdateTime());
        }

        return new StateStatistics(totalCount, categoryCount, maxUpdateTime);
    }
} 